import os
import argparse
import sys

from PIL import Image
import numpy as np

import torch
from torchvision.transforms.functional import to_tensor, to_pil_image

from model import Generator


torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True


def load_image(image_path, x32=False):
    img = Image.open(image_path).convert("RGB")

    if x32:
        def to_32s(x):
            return 256 if x < 256 else x - x % 32
        w, h = img.size
        img = img.resize((to_32s(w), to_32s(h)))
    return img

def test(args):
    device = args.device
    net = Generator()  # 加载预测的模板
    net.load_state_dict(torch.load(args.checkpoint, map_location="cpu"))
    net.to(device).eval()
    print(f"model loaded: {args.checkpoint}")
    num = len(os.listdir("D:/" + args.uid + "/picture"))
    args.output = "D:/" + args.uid + "/picture/output" + str(num+1) + ".jpg"
    image = load_image(args.input, args.x32)
    with torch.no_grad():
        image = to_tensor(image).unsqueeze(0) * 2 - 1
        out = net(image.to(device), args.upsample_align).cpu()
        out = out.squeeze(0).clip(-1, 1) * 0.5 + 0.5
        out = to_pil_image(out)
    out.save(args.output)
    print("D:/" + args.uid + "/picture/output" + str(num+1) + ".jpg")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--uid',
        type=str,
        default='123'
    )
    parser.add_argument(
        '--checkpoint',
        type=str,
        default='E:/Python/FinalAIface/animegan2-pytorch-main/weights/face_paint_512_v2.pt',
    )
    parser.add_argument(
        '--input',
        type=str, 
        default='E:/Python/FinalAIface/animegan2-pytorch-main/samples/inputs/input1.jpg',
    )
    parser.add_argument(
        '--output',
        type=str,
        default='D:/123/picture/output/anime.jpg',
    )
    parser.add_argument(
        '--device',
        type=str,
        default='cpu',
    )
    parser.add_argument(
        '--upsample_align',
        type=bool,
        default=False,
        help="Align corners in decoder upsampling layers"
    )
    parser.add_argument(
        '--x32',
        action="store_true",
        help="Resize images to multiple of 32"
    )
    args = parser.parse_args()
    test(args)
